#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import pandas as pd
from typing import Union, Dict, List
import os
from datetime import datetime
from abc import ABC, abstractmethod

class BaseStrategyExecutor(ABC):
    """
    策略执行器基类，定义通用功能
    """

    def __init__(self, filename: str, initial_capital: float = 200000, trade_volume: int = 400, from_db: bool = False):
        """
        初始化策略执行器

        Args:
            filename: 要读取的CSV文件名
            initial_capital: 初始资金，默认200000
            trade_volume: 每次交易数量，默认400
            from_db: 是否从数据库读取，默认False表示从文件系统读取
        """
        self.filename = filename
        self.from_db = from_db
        self.df: Union[pd.DataFrame, None] = None
        self.signals: List[Dict] = []
        self.initial_capital = initial_capital
        self.current_capital = initial_capital
        self.trade_volume = trade_volume
        self.position = 0
        self.trades: List[Dict] = []
        self.buy_price = 0  # 记录买入价格，用于止损
        self.has_buy = False  # 是否已经买入
        self.last_buy_price = 0  # 上次买入价格
        self.last_buy_has_sell = True  # 上次买入后是否卖出
        self.buy_up_count = 0  # 记录买入赚钱的次数

    def read_csv_content(self) -> Union[pd.DataFrame, None]:
        """
        读取CSV文件内容到pandas DataFrame中

        Returns:
            pd.DataFrame: 读取的CSV数据
            None: 如果读取失败
        """
        try:
            if not self.from_db:
                if not os.path.exists(self.filename):
                    raise FileNotFoundError(f"文件不存在: {self.filename}")

                if not self.filename.lower().endswith('.csv'):
                    raise ValueError("文件必须是CSV格式")

                self.df = pd.read_csv(self.filename, encoding='utf-8')

                # 确保必要的列存在
                required_columns = ['trade_date', 'close']
                missing_columns = [col for col in required_columns if col not in self.df.columns]
                if missing_columns:
                    raise ValueError(f"CSV文件缺少必要的列: {', '.join(missing_columns)}")

                # 确保 trade_date 格式正确并排序
                self.df = self.df.sort_values('trade_date')

                return self.df
            else:
                raise NotImplementedError("数据库读取功能尚未实现")

        except Exception as e:
            raise Exception(f"读取失败: {str(e)}")

    @abstractmethod
    def calculate_moving_averages(self) -> None:
        """
        计算5日、20日和60日移动平均线
        """
        if self.df is None or self.df.empty:
            raise ValueError("DataFrame未初始化或为空")

        # 计算各个周期的移动平均线
        self.df['ma5'] = self.df['close'].rolling(window=5).mean()
        self.df['ma10'] = self.df['close'].rolling(window=10).mean()
        self.df['ma20'] = self.df['close'].rolling(window=20).mean()
        self.df['ma60'] = self.df['close'].rolling(window=60).mean()

    @abstractmethod
    def detect_signals(self) -> None:
        pass

    def execute_trades(self) -> None:
        """
        根据信号执行交易
        """
        for signal in self.signals:
            price = signal['close']
            trade_amount = self.trade_volume

            if signal['signal_type'] == '买入':
                cost = price * trade_amount
                if self.current_capital >= cost:
                    self.current_capital -= cost
                    self.position += trade_amount
                    trade = {
                        'date': signal['trade_date'],
                        'type': '买入',
                        'price': price,
                        'volume': trade_amount,
                        'cost': cost,
                        'position': self.position,
                        'capital': self.current_capital,
                        'ma5': signal['ma5'],
                        'ma10': signal['ma10'],
                        'reason': signal.get('reason', '常规买入')
                    }
                    self.trades.append(trade)

            elif signal['signal_type'] == '卖出':
                if self.position >= trade_amount:
                    revenue = price * trade_amount
                    self.current_capital += revenue
                    self.position -= trade_amount
                    trade = {
                        'date': signal['trade_date'],
                        'type': '卖出',
                        'price': price,
                        'volume': trade_amount,
                        'revenue': revenue,
                        'position': self.position,
                        'capital': self.current_capital,
                        'ma5': signal['ma5'],
                        'ma10': signal['ma10'],
                        'reason': signal.get('reason', '常规卖出')
                    }
                    self.trades.append(trade)

    def calculate_final_assets(self) -> float:
        """
        计算最终资产（现金 + 持仓市值）

        Returns:
            float: 最终总资产
        """
        if not self.df.empty:
            last_price = self.df.iloc[-1]['close']
            position_value = self.position * last_price
            total_assets = self.current_capital + position_value
            return total_assets
        return self.current_capital

    def display_trading_summary(self) -> None:
        """
        显示交易汇总信息
        """
        final_assets = self.calculate_final_assets()
        profit = final_assets - self.initial_capital
        profit_percentage = (profit / self.initial_capital) * 100

        print("\n=== 交易汇总 ===")
        print("-" * 50)

        if self.trades:
            print("\n交易明细:")
            print("-" * 100)
            print(f"{'日期':<12} {'类型':<6} {'价格':>10} {'数量':>8} {'资金变动':>12} {'持仓':>8} {'账户余额':>12} {'ma5':>8} {'ma20':>8} {'ma60':>8}  {'原因':<20}")
            print("-" * 100)

            for trade in self.trades:
                money_change = (-trade['cost'] if trade['type'] == '买入' else trade['revenue'])
                print(f"{trade['date']:<12} "
                      f"{trade['type']:<6} "
                      f"{trade['price']:>10.2f} "
                      f"{trade['volume']:>8d} "
                      f"{money_change:>12.2f} "
                      f"{trade['position']:>8d} "
                      f"{trade['capital']:>12.2f}"
                      f"{trade['ma5']:>12.2f} "
                      f"{trade['capital']:>12.2f} "
                      f"{trade['reason']:<20}")

            print(f"初始资金: {self.initial_capital:,.2f}")
            print(f"当前现金: {self.current_capital:,.2f}")
            print(f"当前持仓: {self.position:,d}股")
            print(f"最终资产: {final_assets:,.2f}")
            print(f"赚钱与亏损比例：{self.buy_up_count} - {len(self.trades) - self.buy_up_count}")
            print(f"总收益: {profit:,.2f} ({profit_percentage:+.2f}%)")
            print("-" * 50)
            print(f"总交易次数: {len(self.trades)}")

    def execute(self) -> None:
        """
        执行完整的数据处理流程
        """
        try:
            # 读取数据
            print(f"正在读取CSV文件: {self.filename}")
            self.read_csv_content()

            if self.df is not None and not self.df.empty:
                # 计算移动平均线
                print("正在计算移动平均线...")
                self.calculate_moving_averages()

                print("正在检测交易信号...")
                self.detect_signals()

                print("正在执行交易模拟...")
                self.execute_trades()

                self.display_trading_summary()

            else:
                print("读取数据失败或数据为空")

        except Exception as e:
            print(f"执行过程中发生错误: {str(e)}")

class FiveTwentyMAStrategy(BaseStrategyExecutor):
    """
    5日均线，10日均线，20日均线均要在60日均线之上
    5日均线大于20日均线时，5日是均线上穿10日均线买入，
    5日均线下穿10日均线卖出
    10% 止盈
    3%止损
    """
    def calculate_moving_averages(self) -> None:
        """
        计算5日、20日和60日移动平均线
        """
        if self.df is None or self.df.empty:
            raise ValueError("DataFrame未初始化或为空")

        # 计算各个周期的移动平均线
        self.df['ma5'] = self.df['close'].rolling(window=5).mean()
        self.df['ma10'] = self.df['close'].rolling(window=10).mean()
        self.df['ma20'] = self.df['close'].rolling(window=20).mean()
        self.df['ma60'] = self.df['close'].rolling(window=60).mean()

    def detect_signals(self) -> None:
        """
        检测均线交叉信号
        规则：
        1. 5日均线和20日均线都在60日均线之上
        2. 检测5日均线上穿或下穿20日均线
        """
        if self.df is None or self.df.empty:
            raise ValueError("DataFrame未初始化或为空")

        # 初始化信号列表
        self.signals = []

        # 创建前一天的移动平均线数据
        self.df['ma5_prev'] = self.df['ma5'].shift(1)
        self.df['ma10_prev'] = self.df['ma10'].shift(1)
        self.df['ma20_prev'] = self.df['ma20'].shift(1)

        # 遍历数据（从第61行开始，确保有足够的数据计算60日均线）
        for i in range(60, len(self.df)):
            current_row = self.df.iloc[i]

            # 检查5日均线和20日均线是否都在60日均线之上
            if current_row['ma5'] > current_row['ma60'] and current_row['ma10'] > current_row['ma60']:
                # 检查是否发生上穿（金叉）
                if (current_row['ma5'] > current_row['ma20'] and
                    self.has_buy == False and current_row['ma5'] > current_row['ma10'] and
                    self.last_buy_has_sell):
                    self.has_buy = True
                    signal = {
                        'trade_date': current_row['trade_date'],
                        'signal_type': '买入',
                        'close': current_row['close'],
                        'ma5': current_row['ma5'],
                        'ma10': current_row['ma10'],
                        'ma20': current_row['ma20'],
                        'ma60': current_row['ma60']
                    }
                    self.last_buy_price = current_row['close']
                    self.last_buy_has_sell = False
                    self.signals.append(signal)

                # 检查是否发生下穿（死叉）
                elif (current_row['ma5'] <= current_row['ma10'] and
                      self.df.iloc[i-1]['ma5'] > self.df.iloc[i-1]['ma10']):
                    self.has_buy = False
                    if self.last_buy_has_sell is False:
                        signal = {
                            'trade_date': current_row['trade_date'],
                            'signal_type': '卖出',
                            'close': current_row['close'],
                            'ma5': current_row['ma5'],
                            'ma10': current_row['ma10'],
                            'ma20': current_row['ma20'],
                            'ma60': current_row['ma60']
                        }
                        if current_row['close'] > self.last_buy_price:
                            self.buy_up_count += 1
                        self.signals.append(signal)
                    self.last_buy_has_sell = True

            if self.has_buy and self.last_buy_has_sell is False and \
                current_row['close'] >= self.last_buy_price * 1.10:
                self.has_buy = False
                signal = {
                    'trade_date': current_row['trade_date'],
                    'signal_type': '卖出',
                    'close': current_row['close'],
                    'ma5': current_row['ma5'],
                    'ma10': current_row['ma10'],
                    'ma20': current_row['ma20'],
                    'ma60': current_row['ma60']
                }
                if current_row['close'] > self.last_buy_price:
                    self.buy_up_count += 1
                self.signals.append(signal)

            if self.has_buy and self.last_buy_has_sell is False and \
                current_row['close'] <= self.last_buy_price * 0.97:
                self.has_buy = False
                signal = {
                    'trade_date': current_row['trade_date'],
                    'signal_type': '卖出',
                    'close': self.last_buy_price * 0.97,
                    'ma5': current_row['ma5'],
                    'ma10': current_row['ma10'],
                    'ma20': current_row['ma20'],
                    'ma60': current_row['ma60']
                }
                if current_row['close'] > self.last_buy_price:
                    self.buy_up_count += 1
                self.signals.append(signal)

class ThreeDayMACrossStrategy(BaseStrategyExecutor):
    """
    连续两天 20日均线上涨
    出现5日均线大于10日均线买入，5日均线小于10日均线卖出
    同时设置3%止损的策略
    """

    def calculate_moving_averages(self) -> None:
        if self.df is None or self.df.empty:
            raise ValueError("DataFrame未初始化或为空")

        # 计算5日和10日移动平均线
        self.df['ma5'] = self.df['close'].rolling(window=5).mean()
        self.df['ma10'] = self.df['close'].rolling(window=10).mean()
        self.df['ma20'] = self.df['close'].rolling(window=20).mean()
        self.df['ma60'] = self.df['close'].rolling(window=60).mean()

        # 计算MA5 > MA10的条件
        self.df['ma5_gt_ma10'] = self.df['ma5'] > self.df['ma10']

        # 计算MA20的上升趋势
        self.df['ma20_prev1'] = self.df['ma20'].shift(1)
        self.df['ma20_prev2'] = self.df['ma20'].shift(2)
        self.df['ma20_prev3'] = self.df['ma20'].shift(3)

        # 判断MA20是否连续3天上升
        self.df['ma20_rising'] = (
            (self.df['ma20'] > self.df['ma20_prev1']) &
            (self.df['ma20_prev1'] > self.df['ma20_prev2']) &
            (self.df['ma20_prev2'] > self.df['ma20_prev3'])
        )

    def detect_signals(self) -> None:
        if self.df is None or self.df.empty:
            raise ValueError("DataFrame未初始化或为空")

        self.signals = []
        consecutive_days = 0
        in_position = False

        for i in range(20, len(self.df)):  # 从第20天开始，确保有足够的数据计算均线
            current_row = self.df.iloc[i]

            # 检查是否需要止损
            if in_position and self.buy_price > 0:
                stop_loss_price = self.buy_price * 0.97  # 3%止损线
                if current_row['close'] <= stop_loss_price:
                    signal = {
                        'trade_date': current_row['trade_date'],
                        'signal_type': '卖出',
                        'close': stop_loss_price,
                        'ma5': current_row['ma5'],
                        'ma10': current_row['ma10'],
                        'ma20': current_row['ma20'],
                        'ma60': current_row['ma60'],
                        'reason': '止损卖出'
                    }
                    self.buy_price = 0
                    if current_row['close'] > self.last_buy_price:
                        self.buy_up_count += 1
                    self.signals.append(signal)
                    in_position = False
                    consecutive_days = 0
                    continue

            # 检查买入条件
            if not in_position:
                if current_row['ma5_gt_ma10'] and current_row['ma20_rising']:
                    consecutive_days += 1
                    if consecutive_days >= 2:
                        signal = {
                            'trade_date': current_row['trade_date'],
                            'signal_type': '买入',
                            'close': current_row['close'],
                            'ma5': current_row['ma5'],
                            'ma10': current_row['ma10'],
                            'ma20': current_row['ma20'],
                            'ma60': current_row['ma60'],
                            'reason': '连续3天MA5>MA10'
                        }
                        self.buy_price = current_row['close']
                        self.last_buy_price = current_row['close']
                        self.signals.append(signal)
                        in_position = True
                else:
                    consecutive_days = 0

            # 检查卖出条件
            elif in_position and not current_row['ma5_gt_ma10']:
                signal = {
                    'trade_date': current_row['trade_date'],
                    'signal_type': '卖出',
                    'close': current_row['close'],
                    'ma5': current_row['ma5'],
                    'ma10': current_row['ma10'],
                    'ma20': current_row['ma20'],
                    'ma60': current_row['ma60'],
                    'reason': 'MA5<MA10'
                }
                self.buy_price = 0
                if current_row['close'] > self.last_buy_price:
                    self.buy_up_count += 1
                self.signals.append(signal)
                in_position = False
                consecutive_days = 0

def main():
    parser = argparse.ArgumentParser(description="股票均线交叉策略分析工具")
    parser.add_argument("--filename", help="要读取的CSV文件名")
    parser.add_argument("--strategy", choices=['three_day_ma', 'five_twenty_ma', 'open_low_diff'], default='three_day_ma',
                      help="选择要使用的策略")
    parser.add_argument("--initial-capital", type=float, default=200000,
                      help="初始资金，默认200000")
    parser.add_argument("--trade-volume", type=int, default=400,
                      help="每次交易数量，默认400")
    parser.add_argument("--from_db", action="store_true",
                      help="是否从数据库读取，默认为False")
    args = parser.parse_args()

    if args.strategy == 'three_day_ma':
        strategy = ThreeDayMACrossStrategy(
            args.filename,
            initial_capital=args.initial_capital,
            trade_volume=args.trade_volume,
            from_db=args.from_db
        )
        strategy.execute()
    elif args.strategy == 'five_twenty_ma':
        executor = FiveTwentyMAStrategy(
            args.filename,
            initial_capital=args.initial_capital,
            trade_volume=args.trade_volume,
            from_db=args.from_db
        )
        executor.execute()
    elif args.strategy == 'open_low_diff':
        executor = OpenLowDiffStrategy(
            args.filename,
            initial_capital=args.initial_capital,
            trade_volume=args.trade_volume,
            from_db=args.from_db
        )
        executor.execute()

class OpenLowDiffStrategy(BaseStrategyExecutor):
    """
    开盘价与最低价的差值大于等于收盘价与开盘价的差值时，
    当天以收盘价买入，第二天以收盘价卖出的策略
    其实就是找长下影线，但是操作的是下影线与开盘价收盘价差不多大小
    """

    def __init__(self, filename: str, initial_capital: float = 200000, trade_volume: int = 400, from_db: bool = False):
        super().__init__(filename, initial_capital, trade_volume, from_db)
        self.is_decrease = None
        self.is_too_high = None
        self.max_day = False

    def calculate_moving_averages(self) -> None:
        """
        由于此策略不使用移动平均线，但为满足基类要求，保留此方法
        """
        if self.df is None or self.df.empty:
            raise ValueError("DataFrame未初始化或为空")

        # 计算移动平均线，用于展示用途
        self.df['ma5'] = self.df['close'].rolling(window=5).mean()
        self.df['ma10'] = self.df['close'].rolling(window=10).mean()
        self.df['ma20'] = self.df['close'].rolling(window=20).mean()
        self.df['ma60'] = self.df['close'].rolling(window=60).mean()

    def detect_signals(self) -> None:
        """
        检测交易信号：
        1. 当天开盘价与最低价的差值 >= 收盘价与开盘价的差值时买入
        2. 买入后第二天收盘时卖出
        """
        if self.df is None or self.df.empty:
            raise ValueError("DataFrame未初始化或为空")

        self.signals = []
        in_position = False

        # 遍历数据，但留出最后一天，因为需要第二天才能卖出
        for i in range(len(self.df) - 1):
            current_row = self.df.iloc[i]

            # 上阴线必须小于下阴线
            # 收盘价比 5日均线高 3% 就不要买了
            if current_row['open'] > current_row['close']:
                # 计算差值
                open_low_diff = abs(current_row['close'] - current_row['low'])
                close_open_diff = abs(current_row['close'] - current_row['open'])
                high_diff = abs(current_row['high'] - current_row['open'])
            else:
                # 计算差值
                open_low_diff = abs(current_row['low'] - current_row['open'])
                close_open_diff = abs(current_row['close'] - current_row['open'])
                high_diff = abs(current_row['high'] - current_row['close'])

            self.is_too_high = current_row['close'] <= 1.02 * current_row['ma5']
            self.is_decrease = current_row['close'] < current_row['open']

            if in_position:
                self.max_day += 1
            # 如果没有持仓，检查买入条件
            if not in_position and open_low_diff >= close_open_diff*2 and high_diff < open_low_diff and \
                self.is_too_high:
                signal = {
                    'trade_date': current_row['trade_date'],
                    'signal_type': '买入',
                    'close': current_row['close'],
                    'ma5': current_row['ma5'],
                    'ma10': current_row['ma10'],
                    'ma20': current_row['ma20'],
                    'ma60': current_row['ma60'],
                    'reason': 'open-low差值大于close-open差值'
                }
                self.max_day = False
                self.signals.append(signal)
                self.last_buy_price = current_row['close']
                in_position = True

            # 如果持仓，第二天卖出
            # 直到股价下跌才卖
            elif in_position and (self.is_decrease or self.max_day >= 1):
                signal = {
                    'trade_date': current_row['trade_date'],
                    'signal_type': '卖出',
                    'close': current_row['close'],
                    'ma5': current_row['ma5'],
                    'ma10': current_row['ma10'],
                    'ma20': current_row['ma20'],
                    'ma60': current_row['ma60'],
                    'reason': '次日卖出'
                }
                self.is_decrease = False
                if current_row['close'] > self.last_buy_price:
                    self.buy_up_count += 1
                self.signals.append(signal)
                in_position = False

if __name__ == "__main__":
    main()
